import json
from typing import List, Tuple, Any, Optional

from langchain_core.documents import Document
import requests
from loguru import logger


class ReRanker:
    def __init__(self, *args: Any, **kwargs: Any):
        self.online_rerank_url = "http://60.205.246.174:8100/rerank"

    def rerank(self, query: str, passages: List[Document], top_k: int = 5) -> List[Tuple[Document, float]]:
        try:
            pairs = [(query, doc.page_content) for doc in passages]
            response = requests.post(self.online_rerank_url, json={"pairs": pairs, "top_k": top_k}).json()
            return [(doc, score) for doc, score in zip(passages, response["scores"])]
        except Exception as e:
            logger.error(f"Rerank failed: {e}")
            raise


class ReRankerServer:
    def __init__(self, *args: Any, **kwargs: Any):
        self.online_rerank_url = "http://10.168.20.20:5002/v1/rerank"
        self.API_KEY = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiI4ZDJkZDdiM2FhZDE0ZDU5YjI0OTZjMzc0MTNiMjJkZiIsIm1ldGhvZHMiOiJQT1NUIiwiaWF0IjoiMTc0OTE3ODg1NiIsImV4cCI6IjE3NDkxODI0NTYifQ.jkCw-LAHHsMFucx30szaFqYdsbI4lzj9BHF6_Lc9el-Elxzy4IYPZIaAnRb6Zb9U38EmshSyiT4SigwOj8UroNcB7bQKAu4ytLTUPQUEyjKPnAgbcs9v9GLwq8OmFbD1ysbKW61G8hblnOa3fjXZbtLHhA5XgZHKsrMKeynI1xXEixswSKij7j-RRRZkITzSQcVWl6a-5l3aqrcZP8VBBrgEHviu2UcgUAKtNXQTsAqLU5ZQHwJ2G4YdlO6ejiLWEUxklFt-pk2kY-uFUoL4qi4f8agIsGrEqrBGtJIdTs4TpD4pQ5oFSWXryduhhSW5VDD2pDYqotAv1T-PqPUDDQ"

    def rerank(self, query: str, passages: List[Document], top_k: int = 5) -> List[Tuple[Document, float]]:
        try:
            headers = {
                "Authorization": f"Bearer {self.API_KEY}",
                "Content-Type": "application/json"
            }
            payload = {
                "model": "bge-reranker-large",
                "query": query,
                "top_n": len(passages),
                "documents": [doc.page_content for doc in passages]
            }
            response = requests.post(self.online_rerank_url, data=json.dumps(payload), headers=headers).json()
            scores = [result["relevance_score"] for result in response["results"]]

            results = []
            for doc, score in zip(passages, scores):
                results.append((doc, score))
            return results
        except Exception as e:
            logger.error(f"Rerank failed: {e}")
            raise